#!/usr/bin/env python3
import mmap
import os
import struct
import cv2
import numpy as np
from datetime import datetime
from logger import setup_logger

logger = setup_logger()

class SharedMemoryReader:
    """修复后的共享内存读取器"""
    def __init__(self, shm_name="/image_shm", max_size=50 * 1024 * 1024):  # 与写入端保持一致
        self.shm_name = shm_name
        self.max_size = max_size
        self._open_shared_memory()
        self.set_save_directories()
    
    def set_save_directories(self):
        """创建保存图像的目录"""
        self.color_dir = "./data/Color6"
        self.depth_dir = "./data/Depth6"
        
        # 创建目录（如果不存在）
        os.makedirs(self.color_dir, exist_ok=True)
        os.makedirs(self.depth_dir, exist_ok=True)
        
        logger.info(f"共享内存图像保存目录已创建: {self.color_dir}, {self.depth_dir}")
    
    def _open_shared_memory(self):
        """打开共享内存"""
        try:
            self.fd = os.open(self.shm_name, os.O_RDONLY)  # 只读模式
            self.shm = mmap.mmap(self.fd, self.max_size, mmap.MAP_SHARED, mmap.PROT_READ)
        except FileNotFoundError:
            logger.warning(f"共享内存文件 {self.shm_name} 不存在，请先启动写入程序")
            raise
    
    def read_images(self):
        """从共享内存读取图像 - 匹配修复后的写入端"""
        try:
            # 新的内存布局：
            # [0-4]: "VALID" 标志 (5字节)
            # [5-12]: timestamp (8字节)
            # [13-20]: color_size + depth_size (8字节)
            # [21+]: 图像数据
            
            self.shm.seek(0)
            
            # 检查有效标志（前5字节）
            flag = self.shm.read(5)
            if flag != b'VALID':
                return None, None, None
            
            # 读取时间戳 (8字节)
            timestamp_data = self.shm.read(8)
            if len(timestamp_data) < 8:
                logger.warning(f"时间戳数据不足: 期望8字节，实际{len(timestamp_data)}字节")
                return None, None, None
            
            timestamp = struct.unpack('Q', timestamp_data)[0]
            
            # 读取大小信息 (8字节)
            size_data = self.shm.read(8)
            if len(size_data) < 8:
                logger.warning(f"大小数据不足: 期望8字节，实际{len(size_data)}字节")
                return None, None, None
                
            color_size, depth_size = struct.unpack('II', size_data)
            
            # print(f"调试: timestamp={timestamp}, color_size={color_size}, depth_size={depth_size}")
            
            # 验证大小是否合理
            if color_size <= 0 or depth_size <= 0 or color_size > 10*1024*1024 or depth_size > 10*1024*1024:
                logger.warning(f"图像大小异常: color_size={color_size}, depth_size={depth_size}")
                return None, None, None
            
            # 读取图像数据（从第21字节开始）
            
            # 读取图像数据
            color_data = self.shm.read(color_size)
            depth_data = self.shm.read(depth_size)
            
            if len(color_data) != color_size or len(depth_data) != depth_size:
                logger.warning(f"数据大小不匹配: 期望彩色={color_size}, 实际={len(color_data)}, 期望深度={depth_size}, 实际={len(depth_data)}")
                return None, None, None
            
            # 解码图像
            color_nparr = np.frombuffer(color_data, dtype=np.uint8)
            depth_nparr = np.frombuffer(depth_data, dtype=np.uint8)
            
            color_img = cv2.imdecode(color_nparr, cv2.IMREAD_COLOR)
            depth_img = cv2.imdecode(depth_nparr, cv2.IMREAD_UNCHANGED)
            
            if color_img is None or depth_img is None:
                logger.warning("图像解码失败")
                return None, None, None
            
            # 转换时间戳为可读格式
            readable_timestamp = datetime.fromtimestamp(timestamp / 1000.0).strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
            
            return readable_timestamp, timestamp, color_img, depth_img
            
        except struct.error as e:
            logger.warning(f"结构解析错误: {e}")
            return None, None, None, None
        except Exception as e:
            logger.warning(f"读取共享内存错误: {e}")
            return None, None, None, None
    
    def save_images(self, timestamp, color_img, depth_img):
        """保存图像到指定目录"""
        try:
            # 使用时间戳作为文件名
            filename = f"{timestamp}"
            
            # 保存彩色图像
            color_path = os.path.join(self.color_dir, f"{filename}.jpg")
            cv2.imwrite(color_path, color_img)
            
            # 保存深度图像
            depth_path = os.path.join(self.depth_dir, f"{filename}.png")
            cv2.imwrite(depth_path, depth_img)
            
            # 检查文件夹中的图片数量
            self._cleanup_old_images(self.color_dir)
            self._cleanup_old_images(self.depth_dir)
            
            return True, color_path, depth_path
            
        except Exception as e:
            logger.warning(f"保存图像错误: {e}")
            return False, None, None
    
    def _cleanup_old_images(self, directory):
        """清理文件夹中的旧图片，确保不超过100张"""
        try:
            # 获取目录中的所有文件
            files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
            
            # 如果文件数量超过100，删除最早的50张
            if len(files) > 100:
                # 按创建时间排序（最早的在前）
                files.sort(key=lambda x: os.path.getctime(os.path.join(directory, x)))
                
                # 删除最早的50张
                for file_to_delete in files[:50]:
                    os.remove(os.path.join(directory, file_to_delete))
                
                logger.info(f"已清理 {directory} 中的50张旧图片")
        except Exception as e:
            logger.warning(f"清理旧图片错误: {e}")
    
    def close(self):
        """关闭共享内存"""
        try:
            self.shm.close()
            os.close(self.fd)
        except:
            pass